import os

from gxl_ai_utils.utils import utils_file


def get_total_scp():
    """"""
    root_dir = "/home/work_nfs10/data/aslp_spider_data_k2_fbank/pici_45093H"
    main_dataset_dir_name_list, _ = utils_file.do_listdir(root_dir, return_path=False)
    main_dataset_dir_name_list = [item for item in main_dataset_dir_name_list if item.startswith('main_5_')]
    print(main_dataset_dir_name_list)
    total_wav_dict = {}
    total_text_dict = {}
    output_dir = "/home/work_nfs10/data/aslp_spider_data_k2_fbank/pici_45093H/main_6_res_10219H"
    output_wav_scp_path = os.path.join(output_dir, 'wav.scp')
    output_txt_scp_path = os.path.join(output_dir, 'text')

    for main_dataset_dir_name in main_dataset_dir_name_list:
        print(f'handling {main_dataset_dir_name}')
        root_dir_partition = os.path.join(root_dir, main_dataset_dir_name)
        nodes_list, _ = utils_file.do_listdir(root_dir_partition, return_path=False)
        print(f'{nodes_list}')
        for node in nodes_list:
            print(f'handling {node}')
            node_dir = os.path.join(root_dir_partition, node, 'split')
            split_names ,_ =  utils_file.do_listdir(node_dir, return_path=False)
            print(split_names)
            for split_name in split_names:
                print(f'handling {split_name}')
                split_dir = os.path.join(node_dir, split_name)
                fbank_file_path1 = os.path.join(split_dir, 'fbank', "aslp-spider-data_cuts_train.jsonl.gz")
                fbank_file_path2 = os.path.join(split_dir, 'fbank', "aslp-spider-data_cuts_train.jsonl")
                if os.path.exists(fbank_file_path1) or os.path.exists(fbank_file_path2):
                    continue
                wav_scp_path = os.path.join(split_dir, 'wav.scp')
                text_scp_path = os.path.join(split_dir, 'text')
                if not os.path.exists(wav_scp_path) or not os.path.exists(text_scp_path):
                    continue
                wav_dict_i = utils_file.load_dict_from_scp(wav_scp_path)
                text_dict_i = utils_file.load_dict_from_scp(text_scp_path)
                total_wav_dict.update(wav_dict_i)
                total_text_dict.update(text_dict_i)

    utils_file.write_dict_to_scp(total_wav_dict, output_wav_scp_path)
    utils_file.write_dict_to_scp(total_text_dict, output_txt_scp_path)

if __name__ == '__main__':
    get_total_scp()
